import os
import cv2
import numpy as np
import onnxruntime
import time

class YOLOV5():
    def __init__(self,onnxpath):
        self.onnx_session=onnxruntime.InferenceSession(onnxpath)
        self.input_name=self.get_input_name()
        self.output_name=self.get_output_name()
    #-------------------------------------------------------
	#   获取输入输出的名字
	#-------------------------------------------------------
    def get_input_name(self):
        input_name=[]
        for node in self.onnx_session.get_inputs():
            input_name.append(node.name)
        return input_name
    def get_output_name(self):
        output_name=[]
        for node in self.onnx_session.get_outputs():
            output_name.append(node.name)
        return output_name
    #-------------------------------------------------------
	#   输入图像
	#-------------------------------------------------------
    def get_input_feed(self,img_tensor):
        input_feed={}
        for name in self.input_name:
            input_feed[name]=img_tensor
        return input_feed
    #-------------------------------------------------------
	#   1.cv2读取图像并resize
	#	2.图像转BGR2RGB和HWC2CHW
	#	3.图像归一化
	#	4.图像增加维度
	#	5.onnx_session 推理
	#-------------------------------------------------------
    def inference(self,img):
        # 图像预处理为等比缩放然后加黑色边
        img = cv2.resize(img, (640, 360))
        or_img = cv2.copyMakeBorder(img, 140, 140, 0, 0, cv2.BORDER_CONSTANT)
        # 通道转换并归一化
        img=or_img[:,:,::-1].transpose(2,0,1)  #BGR2RGB和HWC2CHW
        img=img.astype(dtype=np.float32)
        img/=255.0
        img=np.expand_dims(img,axis=0)
        input_feed=self.get_input_feed(img)
        pred=self.onnx_session.run(None,input_feed)
        return pred,or_img
    
def restore_image(image_640x640, original_size=(1280, 720)):
    """
    将预处理后的 640×640 图像（包含黑边）还原为原始 1280×720 图像。

    参数：
        image_640x640: 预处理图像，shape (640, 640, 3)
        original_size: 原图尺寸 (W, H)，例如 (1280, 720)
    
    返回：
        还原后的图像，shape (720, 1280, 3)
    """
    # 1. 去除上下黑边 ⇒ 保留中间 360 像素
    cropped = image_640x640[140:500, :]  # shape (360, 640, 3)

    # 2. resize 回原始尺寸
    restored = cv2.resize(cropped, original_size, interpolation=cv2.INTER_LINEAR)

    return restored